# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         get_features.py
# Date:           2025-05-24
# Description:    本模块负责从图像中提取骰子目标的关键特征，包括：
#                 - 使用 YOLO 进行目标检测与分类；
#                 - 基于颜色分割与轮廓分析提取骰子表面角点；
#                 - 结合 HSV 颜色空间判断提升识别鲁棒性；
#                 - 整合 PnP、BA、PoseNet 和融合算法进行姿态估计；
#                 整个骰子姿态估计系统的核心特征提取组件。
# ==============================================================================

import cv2
import numpy as np

from utils.grabscreen import grab_screen, setup_window
from core.identify_surface import hsv_to_color_name, get_mode_channel, color_split, extract_corner_points
from configs.config import Config
from configs.norm_data import colors_dict


# 主要处理图像过程
def process_frame(img, display_options, control_params):
    # 使用YOLO进行检测
    result = Config.YOLO_MODEL(img, verbose=False)[0]

    masks = result.boxes.conf >= control_params['yolo_conf']
    boxes = result.boxes[masks]
    confidences = boxes.conf.tolist()
    class_indices = boxes.cls.int().tolist()
    draw_img = img.copy()

    mask_list = []
    corners_list = []
    yolo_datas = []
    found_flag = False

    # 遍历每个检测到的物体
    for box, conf, cls_idx in zip(boxes, confidences, class_indices):
        color_name = colors_dict[cls_idx]

        x, y, w, h = map(int, box.xywh[0])  # 获取中心坐标 + 宽高
        offset_x1 = max(0, int(x - w // 2) - control_params['box_offset_value'])
        offset_y1 = max(0, int(y - h // 2) - control_params['box_offset_value'])
        offset_x2 = min(img.shape[1], int(x + w // 2) + control_params['box_offset_value'])
        offset_y2 = min(img.shape[0], int(y + h // 2) + control_params['box_offset_value'])

        # 裁剪检测区域
        boxes_roi = img[offset_y1:offset_y2, offset_x1:offset_x2]
        color_section = color_split(boxes_roi, color_name)

        section_contours, section_hierarchy = cv2.findContours(color_section, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        for index, contour in enumerate(section_contours):
            if section_hierarchy[0][index][3] == -1 and section_hierarchy[0][index][2] >= 0 and cv2.contourArea(contour) > control_params['contour_area']:
                found_flag = True

                mask = np.zeros(color_section.shape[:2], dtype=np.uint8)
                cv2.drawContours(mask, [contour], -1, 255, thickness=cv2.FILLED)
                mask_list.append(mask)
                # cv2.imshow("mask", mask)
                roi_masked = cv2.bitwise_and(boxes_roi, boxes_roi, mask=mask)
                hue, sat, val = get_mode_channel(roi_masked)
                # 此处再加条件判断可去除一定误识别
                if control_params['use_hsv_flag']:
                    hsv_color_name = hsv_to_color_name(hue, sat, val)

                cnt_translated = contour + (offset_x1, offset_y1)
                corners = extract_corner_points(cnt_translated)
                corners_list.append(corners)
                if display_options['show_yolo']:
                    yolo_datas.append([cls_idx, conf, offset_x1, offset_y1, offset_x2, offset_y2])

    if found_flag:
        draw_img = handle_corners(draw_img, mask_list, corners_list, yolo_datas, display_options, control_params)
        # cv2.imshow("draw_img", draw_img)
        # cv2.waitKey(0)
        return draw_img
    else:
        return draw_img



# 处理角点并进行 姿态估计
def handle_corners(draw_img, mask_list, corners_list, yolo_datas, display_options, control_params):
    from utils.utils import find_best_pose
    from utils.visualization import compute_pnp_results, compute_ba_results, compute_posenet_results, compute_fusion_result, draw_projections,  display_pose_info, draw_error_text, draw_pose_text, draw_yolo_datas

    # Step 1: PnP
    rvec_list, tvec_list, pnp_imgpts_list, pnp_errors = compute_pnp_results(
        corners_list
    )

    # Step 2: BA
    ba_rvecs, ba_tvecs, ba_imgpts_list, ba_errors = compute_ba_results(
        corners_list, rvec_list, tvec_list, control_params['huber_loss']
    )

    # Step 3: PoseNet
    posenet_rvecs, posenet_imgpts_list, posenet_errors = compute_posenet_results(
        corners_list, mask_list, tvec_list
    )

    # Step 4: Fusion
    fusion_imgpts, fusion_error = compute_fusion_result(
        corners_list, rvec_list, tvec_list, ba_rvecs, ba_tvecs, posenet_rvecs
    )

    # Step 5: Draw projections
    texts = draw_projections(
        draw_img, pnp_imgpts_list, ba_imgpts_list, posenet_imgpts_list, fusion_imgpts,
        display_options, control_params, pnp_errors, ba_errors, posenet_errors, fusion_error
    )

    # Step 6: Draw text
    draw_error_text(draw_img, texts, control_params['text_font_scale'], control_params['text_font_thickness'])

    # Step 7: Draw yolo datas
    if display_options['show_yolo']:
        draw_yolo_datas(draw_img, yolo_datas, control_params['yolo_font_scale'], control_params['normal_line_thickness'])


    # Step 8: Get best pose
    best_method, best_idx, best_rvec, best_tvec, min_error = find_best_pose(corners_list, pnp_errors, ba_errors, posenet_errors,
                                                                            rvec_list, tvec_list, ba_rvecs, ba_tvecs, posenet_rvecs)


    # Step 9: Get display texts
    texts_to_display = display_pose_info(best_method, best_rvec, best_tvec, min_error)

    # Step 10: Draw pose text
    draw_pose_text(draw_img, texts_to_display, control_params['text_font_scale'], control_params['text_font_thickness'])

    return draw_img

